package com.ljt.aicodehelp.ai.guardrail;

import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.guardrail.InputGuardrail;
import dev.langchain4j.guardrail.InputGuardrailResult;

import java.util.Set;

// 敏感词拦截器
public class SensitiveWordGuardRail implements InputGuardrail {
    private static Set<String> sensitiveWordSet = Set.of("法西斯万岁");

    @Override
    public InputGuardrailResult validate(UserMessage userMessage) {
        var inputText = userMessage.singleText();

        for (var s : sensitiveWordSet) {
            if (inputText.contains(s)) {
                return fatal("有敏感词：" + s);
            }
        }

        return success();
    }
}
